feat(dynacell): 2D in-focus VS model track - #487

…D evaluation (#484)

* feat(dynacell/eval): support 2D pixel metrics (FRC) alongside 3D (FSC)

Make the pixel-metric path dimension-flexible so the same code scores a
single in-focus 2D projection as well as the full 3D volume, following
cubic's convention of dispatching on input rank:

- ssim() accepts 2-D (H, W) or 3-D (D, H, W) input and derives
  spatial_dims from ndim, instead of hard-rejecting non-3D.
- compute_pixel_metrics() uses ring-based frc_resolution (FRC_Resolution)
  for 2D and keeps shell-based fsc_resolution (*_FSC_Resolution) for 3D,
  and slices spacing to the trailing spatial axes. PCC/NRMSE/PSNR/
  Spectral_PCC are shape-agnostic. The 3D path is byte-identical.

This is the metric-side enabler for adding in-focus 2D columns to the
benchmark without replacing the existing full-3D pixel metrics.

The test change also fixes a pre-existing stale stub harness: the cubic
stubs lacked glcm_features / cubic.scipy / cubic.skimage, so the guarded
top-level import fell back to all-None and every stubbed test failed at
HEAD. Adds those stubs plus 2D/3D dispatch, FRC-vs-FSC, and real-cubic
integration tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011qMvW4Um7irENLqZ1XnZY2

* feat(dynacell/eval): add io.exclude_fov_names to skip incomplete-prediction FOVs

evaluate_predictions can now drop specific FOVs (matched by full name
0/0/fov0011 or leaf fov0011) via io.exclude_fov_names, applied uniformly to
pred/gt/seg positions so the strict position-count alignment still holds.
Unlike limit_positions it does not set partial_walk, so deep-feature caches
self-heal normally. Eval-side analog of the predict-side exclude_fov_names
(--resume-predict): evaluate the finished FOVs of a partially-run predict.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* test(dynacell/eval): cover io.exclude_fov_names FOV exclusion

End-to-end test driving evaluate_predictions on the 3-position CPU fixture:
excluding one position by its full name (A/1/1) or its leaf (1) leaves
N_POSITIONS-1 positions' worth of rows and keeps the strict pred/gt
count-alignment. Mirrors the limit_positions integration test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* feat(dynacell): 2D in-focus VS model track (#487)

* feat(viscy-models): expose downsample_z on Unet3d for Z-preserving 2D use

Unet3d hardcoded downsample_z=True into UNet3DBase, so it could not run
Z-preserving at Z=1 (the base D%2**depth divisor check trips on D=1) and
rejected a downsample_z kwarg with TypeError. Expose downsample_z (default
True, preserving FNet3D) and forward it to the base. With downsample_z=False
the encoder uses stride (1,2,2) / Z-preserving ConvTranspose3d and the base
skips the D divisor check, enabling the FNet2D track to run the FNet
recursive U-Net at Z=1.

Regression tests: default keeps downsamples_z True; downsample_z=False runs
a (B,1,1,64,64) forward preserving Z; and skips the Z divisor check for a
non-divisible Z. 11/11 pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* feat(dynacell): add offline focus-slab store extractor for the 2D track

Pre-extract <name>_focus.zarr for 2D in-focus VS training: per position/T,
estimate the in-focus plane on the phase channel by reusing the eval
estimator (dynacell.evaluation.focus.estimate_focus_plane + the
MIDBAND_FRACTIONS constant, no reimplementation, so train and eval land
on the same plane), then cut a 2*hw+1-plane slab (hw=2 -> 5 planes) of all
channels. A new clamp_shift_slab shifts the whole window inward at edge-focus
FOVs (vs the eval's clipping focus_slab_from_plane) so every position yields
a uniform rectangular depth for SlidingWindowDataset. omero/normalization and
custom zattrs are preserved; a focus_slab_extraction provenance zattr records
the params. pixel_size defaults to the store lateral spacing (--pixel-size
override). zarr v3 / OME-Zarr v0.5, context-managed I/O.

9 integration tests (synthetic OME-Zarr with injected focus planes: shape,
channel/zattr preservation, plane centering, clamp-shift edges, limit/overwrite,
missing-channel raise, CLI) pass.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* feat(dynacell): add 2D-track config overlays (fcmae_vscyto2d + fnet2d)

Overlays for the 2D in-focus VS model track (fit + predict), mirroring the
3D FCMAE/FNet overlays with the 2D deltas:
- recipes/models/fnet2d.yml: architecture FNet3D + model_config.downsample_z:
  false, in_stack_depth 1 (Z-preserving FNet run at Z=1 via the new Unet3d
  downsample_z passthrough).
- data_overlays fcmae_vscyto2d_fit / fnet2d_fit: z_window_size 1; flatten Z
  augmentations (drop Z rotate/scale/sigma_z + the [15,...]/[32,...] z-crop;
  keep XY); RandWeightedCrop / center-crop roi Z -> 1.
- model_overlays fcmae_vscyto2d_fit / fnet2d_fit: 2D model_config (FCMAE stem
  [1,2,2], in_stack_depth 1; FNet2D downsample_z false), single_gpu topology
  (drop the DDP 3h-timeout block), FCMAE lr 4e-4/200 epochs/16-mixed, FNet2D
  200k steps/32-true.
- model_overlays fcmae_vscyto2d_predict / fnet2d_predict: re-declare the full
  2D model_config + data.z_window_size 1 so predict rebuilds a matching model
  before load_state_dict (engine.py:308).

Render-verified via submit_benchmark_job.py --print-resolved-config.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* feat(dynacell): add 2D nucleus/a549 smoke fit+predict leaves (fcmae_vscyto2d)

Phase-B smoke pair for the 2D track: fcmae_vscyto2d_scratch, nucleus, a549.
Fit reads the offline in-focus slab store dual_nucl_memb_all_focus.zarr inline
(bypasses dataset_ref, as A549 leaves do), z_window_size 1, target_channel
Nuclei, RandWeightedCrop re-keyed to Nuclei. Predict runs the 2D model over
the unchanged 3D test store at z_window_size 1 -> full-Z prediction.zarr under
the canonical directory grammar. Predict target_channel Nuclei satisfies the
Nuclei_prediction eval channel contract. ckpt_path/warmup_steps are Phase-B/F
TODOs (render clean; not loaded at compose time).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* fix(dynacell): use identity Z scale in FCMAE-2D affine augmentation

The FCMAE-2D fit data overlay set scale_range Z to [0, 0], intending
"no Z scaling". kornia's RandomAffine3D treats scale as an absolute
factor (1.0 = identity), so 0 zeroes diagonal element 3 of the 4x4
affine matrix; warp_affine3d's torch.linalg.inv then raises a singular-
matrix _LinAlgError on the first GPU-augmented training batch. Use the
identity factor [1, 1] and correct the mapping comment. Affects both
FCMAE-2D leaves (scratch + pretrained), which compose this overlay.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* fix(dynacell): size FCMAE-2D single-GPU batch to fit the 48 GB pool

The FCMAE-2D fit overlay ran single-GPU (2D nets are light) but kept the
3D overlay's batch_size: 32, which there was split across ddp_4gpu
(8/GPU). On one device 32 samples/step OOMs a 48 GB a40/a6000/l40s at
~47 GB, and the long GPU profile sets constraint: null, so any FCMAE-2D
fit could land there and die on the first step. Lower to 16 (a verified
bs=8 uses ~half the 48 GB, so 16 keeps ~2x headroom) so the fit stays on
the cheap pool instead of being forced onto scarce 80 GB Hopper GPUs
across all 36 planned fits. Effective batch is intentionally not matched
to the 3D 4-GPU global batch; the 2D nets are tuned on their own.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* feat(dynacell): add fcmae_vscyto2d_scratch fit leaves for the 2D track

Author the remaining 11 UNeXt2-2D (FCMAE, pretraining=False, random init)
fit leaves — 4 organelles x 3 train_sets, minus the pre-existing
nucleus/a549_mantis smoke-validated exemplar. Each mirrors the 3D
fcmae_vscyto3d_scratch leaf with 2D geometry: composes the
fcmae_vscyto2d fit overlays (in_stack_depth=1, stem [1,2,2],
z_window_size=1), single-GPU topology, and reads the offline 5-plane
in-focus _focus.zarr slab stores.

iPSC single-set leaves set benchmark.dataset_ref: null and inline
source_channel/target_channel/data_path, because the _compose_hook
resolver raises when a leaf co-declares data_path alongside a full
inherited dataset_ref. Joint leaves inline a BatchedConcatDataModule
over the ipsc + a549 focus stores with batch_size 4 (joint mode does
not divide by num_samples, so 4 x 4 = 16 samples/step matches the
single-set 16) and the Z-flattened augmentation policy.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* feat(dynacell): add fcmae_vscyto2d_pretrained fit leaves for the 2D track

Author the 12 VSCyto2D (FCMAE, pretraining=False, FCMAE-pretrained encoder
init) fit leaves — 4 organelles x 3 train_sets. Identical to the
fcmae_vscyto2d_scratch leaves except the model block loads the 2D FCMAE
encoder (encoder_only: true, ckpt from the fcmae-2d
hek-a549-bj5a-livecell-ctmc-400ep pretrain) and finetunes at lr 0.0002,
mirroring the cytoland VSCyto2D recipe. Same 2D geometry, focus stores,
iPSC dataset_ref: null handling, and joint BatchedConcatDataModule as the
scratch counterpart. er/mito keep the 3D source's _ws8500 wandb-name
suffix so 2D and 3D pretrained runs stay greppable together.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* feat(dynacell): add fnet2d fit leaves for the 2D track

Author the 12 FNet2D fit leaves — 4 organelles x 3 train_sets. Each
mirrors the 3D fnet3d_paper leaf but composes the fnet2d fit overlays
(FNet3D architecture with downsample_z: false + in_stack_depth 1 =
Z-preserving 2D net, z_window_size 1, [1,64,64] crops), keeps the
single-GPU topology, and reads the offline _focus.zarr slab stores.
Retains FNet's mean/std Structure normalization and 8-crops-per-FOV
sampling. iPSC single-set leaves use dataset_ref: null + inlined
channels/path (resolver co-declaration guard); joint leaves inline a
BatchedConcatDataModule over the ipsc + a549 focus stores. sbatch mem
drops from the 3D leaves' 512G to 128G since the 5-plane focus stores
are far smaller than the full-Z sources.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* docs: correct Unet3d Z-divisibility docstring for downsample_z

The class docstring unconditionally claimed the model downsamples all
three spatial dimensions and that Z must be divisible by 2**depth. With
the downsample_z parameter, Z is preserved (stride (1,2,2)) when
downsample_z=False, so the Z-divisibility constraint applies only when
downsample_z=True. Also replace the hardcoded "D % 16" (correct only for
depth=4) with the depth-general "D % 2**depth". Docstring only; no
behavior change.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* fix(dynacell): pin OME-Zarr version=0.5 on focus-slab output store

The module docstring states the output is zarr v3 / OME-Zarr v0.5, but
the open_ome_zarr write call relied on iohub's default version. Pin
version="0.5" explicitly so the output format is deterministic regardless
of the iohub default, matching build_cpdino_seg_cleaned.py.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* docs: fix stale bs=32 comments in fcmae_vscyto2d model overlay

The batch size in the sibling data_overlays/fcmae_vscyto2d_fit.yml was
lowered 32->16 for single-GPU OOM headroom, but two comments in the model
overlay still referenced bs=32. Update both to bs=16 to match the actual
batch_size the composed leaves use.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* fix(viscy-models): restore 5D input guard in UNet3DBase.forward

An explicit `if x.ndim != 5` guard once lived in Unet3d.forward but was
dropped in refactor 3bad6cd when that method was removed in favor of the
shared UNet3DBase.forward, which never re-added it. A 4D input (e.g. a
squeezed Z=1 dim) then hits Conv3d and raises a confusing "expected N
channels" error instead of a clear message. The 2D track's Z=1 /
downsample_z=False path newly exercises this, so restore the guard and
document both raises. Add a regression test.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

---------

Co-authored-by: Alexandr Kalinin <alxndrkalinin@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* refactor(dynacell/eval): share focus estimator defaults via focus constants

Hoist the mantis acquisition defaults (na_det=1.35, lambda_ill=0.450)
into module constants DEFAULT_NA_DET / DEFAULT_LAMBDA_ILL in
evaluation/focus.py, use them as the read_focus_compute_config defaults,
and import them into the focus-slab extractor instead of re-declaring the
literals. Single source of truth so the training slab and the eval slab
share the same in-focus plane (train/eval cannot silently diverge).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* refactor(dynacell): simplify focus-slab extractor internals

Tidy ExtractionSummary and the extraction loop without behavior change:
derive n_positions as a property of focus_planes (drop the redundant
field and its manual counter), build the summary once the pixel size is
resolved (drop the float("nan") placeholder + reassignment), inline the
single-use _split_position_name helper into name.split("/"), and unpack
the leading shape dims via data.shape[:3].

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* refactor(dynacell/eval): simplify freq_spacing rank inference

compute_pixel_metrics always receives a (Z, Y, X) spacing tuple, so
reduce it to the input rank unconditionally with list(spacing)[-ndim:]
and drop the dead scalar-spacing guard. Also drop the redundant int()
around pred_xp.ndim (already int). Byte-identical for the 3D path.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* fix(dynacell): train FCMAE-2D in 32-true to stop fp16 NaN divergence

All 24 FCMAE_VSCyto2D fits (scratch + pretrained, every train_set ×
organelle) diverged to NaN loss under precision: 16-mixed — top-k
checkpoints frozen at epoch 0-15, final loss/validate=NaN despite
reaching epoch 199/200. FNet2D (32-true) and the 3D FCMAE (same
16-mixed + MixedLoss) both train fine, so the failure is specific to the
2D FCMAE under fp16.

Verified empirically (scratchpad/repro_fcmae2d_nan.py): MixedLoss /
ms_ssim_25d on a real median/iqr-normalized target with near-untrained
predictions is finite in both fp32 and fp16 autocast, so the NaN is not
the loss forward (nor zero-divisor normalization) but fp16 training
dynamics — activation/gradient overflow the GradScaler misses. 32-true
removes fp16 entirely; 2D nets are light so the cost is negligible.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_013aLSCR6RRyUwjHXzF2RAKg

* fix(viscy-utils): stop bf16 SSIM NaN that diverged 2D FCMAE training

_compute_ssim_and_cs_bf16 runs its mean / mean-of-squares convolutions in
bf16 regardless of trainer precision. On large-magnitude z-scored FCMAE
targets (a channel with a tiny iqr, e.g. nucleus iqr~1.8, normalizes to
values in the hundreds) the variance estimate mu_xx - mu_x**2
catastrophically cancels and rounds *negative*. At depth=1 (2D FCMAE) each
SSIM window has only 1*H_k*W_k samples, so the cancellation is far more
likely than at depth=15 (3D, where the extra Z samples keep the estimate
stable) — a near-flat window then collapses the SSIM denominator to ~0 and
yields intermittent NaN loss that trains fine for a while then diverges.
clamp=True in MixedLoss cannot repair it (clamp bounds magnitude, not NaN),
so the poisoned gradient corrupts every weight.

Fix, in the shared helper:
- clamp both diagonal variances to >=0 (a variance is non-negative by
  definition), flooring the CS denominator at c2;
- bound the covariance by Cauchy-Schwarz (|sigma_xy| <= sqrt(sigma_x*sigma_y)).
  Exact fp32 recovers cs=1 on a flat window because the equally-negative
  sigma_x/sigma_y/sigma_xy cancel in the ratio; clamping only the diagonal
  terms leaves an unbalanced negative sigma_xy and inflates |cs| past 10 on
  background windows (which dominate a VS image). The bound restores cs=1;
- eps-floor both denominators so a zero data_range (c1=c2=0) cannot yield 0/0.

On well-conditioned windows the bound is slack and the clamps are no-ops:
3D correlated-pair aggregate SSIM and gradient-flow cosine vs MONAI fp32 are
unchanged within existing tolerances.

Two deterministic, hardware-independent regression tests are added: zero
data_range stays finite (baseline NaN on every device), and near-flat
large-magnitude windows keep cs~=1 (the unbounded-covariance variant fails
with |cs-1| > 10). The negative-variance NaN itself is a bf16 rounding
knife-edge not unit-reproducible across GPU architectures; it is validated
end-to-end (the exact failing fit ran clean after the clamp).

* fix(dynacell): revert FCMAE-2D fit to 16-mixed now that the SSIM NaN is fixed

The 32-true switch (93990d1) was a wrong hypothesis: the 2D FCMAE NaN was
the bf16 SSIM denominator collapse in ms_ssim_25d, not fp16 training
dynamics. _compute_ssim_and_cs_bf16 runs bf16 convolutions regardless of
trainer precision, so 32-true never addressed the real cause. The preceding
metrics fix clamps the variances and bounds the covariance at the source, so
16-mixed is safe again — restoring the lighter, faster precision the 3D
FCMAE also uses.

---------

Co-authored-by: Alexandr Kalinin <alxndrkalinin@users.noreply.github.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>